from neo4j import GraphDatabase
from embedding import get_embedding


def main():
    driver = GraphDatabase.driver("bolt://localhost:7687", auth=("neo4j", "12345678"))
    # 定义要处理的节点的标签
    node_labels = ["Book", "Author"]
    for node_label in node_labels:
        result = driver.execute_query(f"match (n:{node_label}) return n.name as name")
        nodes = result.records
        print(f"共找到{len(nodes)}个{node_label}节点")
        for record in nodes:
            data = record.data()
            name = data.get("name")
            embedding = get_embedding(name)
            driver.execute_query(
                f"match (n:{node_label}) where n.name = $name set n.embedding=$embedding",
                name=name,
                embedding=embedding,
            )
    print("数据处理完成")
    driver.close()


if __name__ == "__main__":
    main()
